import time

from RelayMemory import ReplayMemory, PriorityReplayMemory
from utilenn import numpy_image_to_tensor_image, tensor_image_to_numpy_image
import gym
import gym_ple

from OptimizeTensor import OTensor


class FlappyEnvironment(object):
    def __init__(self, mem_size=5000):
        self.game = gym.make('FlappyBird-v0')
        self.game.seed(int(time.time()))

        self.action_size = 2
        self.current_state = None
        self.current_frame = None
        self.t = 0

        self.mem = PriorityReplayMemory(mem_size)
        # self.mem = ReplayMemory(mem_size)

    def reset(self):
        self.t = 0
        self.game.reset()
        self.game.step(1)

        self.current_frame = self.get_screen()
        self.current_state = self.get_screen() - self.current_frame

    def get_screen(self):
        return numpy_image_to_tensor_image(self.game.render('rgb_array')).unsqueeze(0)

    def raw_screen(self):
        return self.game.render('rgb_array')

    def step(self, action):
        _state, reward, done, _obs = self.game.step(action)

        skip = 0
        while not done and skip < 2:
            _state, next_reward, done, _obs = self.game.step(1)
            reward += next_reward
            skip += 1

        next_frame = self.get_screen()
        next_state = next_frame - self.current_frame

        if reward != 0:
            print('the reward', reward)

        if done:
            next_state = None

        self.mem.push(self.current_state, OTensor.LongTensor([[action]]), next_state,
                      OTensor.FloatTensor([reward]), done)

        self.current_state = next_state
        self.current_frame = next_frame

        return done


if __name__ == '__main__':
    from matplotlib import pyplot as plt

    env = FlappyEnvironment()
    env.reset()

    env.step(1)

    print(env.current_state.size())

    img = plt.imshow(tensor_image_to_numpy_image(env.current_state), cmap='gray')
    #
    for _ in range(20):
        env.step(0)
        img.set_data(tensor_image_to_numpy_image(env.current_state))
        plt.draw()
        plt.pause(0.1)
    #
    plt.show(block=True)
#
# print(env.get_screen().size())
#
# for i in range(100):
#     if env.step(1):
#         break
#

#
